#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
SepMark 水印解码测试脚本
用于测试已嵌入水印图片的解码功能
"""

import os
import torch
import numpy as np
from PIL import Image
import yaml
import json
from easydict import EasyDict
from torchvision import transforms
from network.Dual_Mark import Network
import argparse
import time
import random


class WatermarkDecoder:
    def __init__(self, config_path=r'D:\deepfake\SepMark-main\SepMark-main\cfg\test_DualMark.yaml'):
        """
        初始化水印解码器

        Args:
            config_path: 配置文件路径
        """
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"使用设备: {self.device}")

        # 加载配置
        self.load_config(config_path)

        # 初始化网络
        self.setup_network()

        # 设置图像转换
        self.transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
        ])

        print("水印解码器初始化完成")

    def load_config(self, config_path):
        """加载配置文件"""
        with open(config_path, 'r') as f:
            test_args = EasyDict(yaml.load(f, Loader=yaml.SafeLoader))

        self.result_folder = "results/" + test_args.result_folder
        self.model_epoch = test_args.model_epoch
        self.strength_factor = test_args.strength_factor

        # 加载训练配置
        with open(r'D:\deepfake\SepMark-main\SepMark-main\cfg\train_DualMark.yaml') as f:
            train_args = EasyDict(yaml.load(f, Loader=yaml.SafeLoader))

        self.message_length = train_args.message_length
        self.message_range = train_args.message_range
        self.attention_encoder = train_args.attention_encoder
        self.attention_decoder = train_args.attention_decoder
        self.weight = train_args.weight
        self.noise_layers_R = train_args.noise_layers.pool_R
        self.noise_layers_F = train_args.noise_layers.pool_F

        print(f"配置加载完成: 模型epoch={self.model_epoch}, 消息长度={self.message_length}")

    def setup_network(self):
        """设置网络结构并加载模型"""
        # 初始化网络
        self.network = Network(
            self.message_length,
            self.noise_layers_R,
            self.noise_layers_F,
            self.device,
            1,  # batch_size = 1
            0.0002,  # lr (推理时不用)
            0.5,  # beta1 (推理时不用)
            self.attention_encoder,
            self.attention_decoder,
            self.weight
        )

        # 加载预训练模型
        model_path = r"D:\deepfake\SepMark-main\SepMark-main\results\Dual_watermark_256_128_0.1_0.0002_0.5_se_se_1_10_10_10_0.1_2023_04_18_16_29_54\models\EC_90.pth"
        if os.path.exists(model_path):
            self.network.load_model_ed(model_path)
            print(f"模型加载成功: {model_path}")
        else:
            print(f"⚠️  警告: 模型文件未找到: {model_path}")
            print("⚠️  使用随机初始化的网络进行测试")

        # 设置为评估模式
        self.network.encoder_decoder.eval()

    def load_watermark_info(self, watermark_info_path):
        """
        加载水印信息文件

        Args:
            watermark_info_path: 水印信息文件路径

        Returns:
            水印信息字典
        """
        if not os.path.exists(watermark_info_path):
            raise FileNotFoundError(f"水印信息文件未找到: {watermark_info_path}")

        with open(watermark_info_path, 'r') as f:
            watermark_info = json.load(f)

        print(f"水印信息加载成功:")
        print(f"  - 水印种子: {watermark_info['watermark_seed']}")
        print(f"  - 消息长度: {watermark_info['message_length']}")
        print(f"  - 消息范围: {watermark_info['message_range']}")
        print(f"  - 强度因子: {watermark_info['strength_factor']}")
        print(f"  - 处理图片数: {watermark_info['successful_count']}")

        return watermark_info

    def generate_original_message(self, seed):
        """
        根据种子重新生成原始水印消息

        Args:
            seed: 水印种子

        Returns:
            原始水印消息tensor
        """
        np.random.seed(seed)
        message = torch.Tensor(
            np.random.choice([-self.message_range, self.message_range],
                             (1, self.message_length))
        ).to(self.device)
        return message

    def load_watermarked_image(self, image_path):
        """
        加载水印图像

        Args:
            image_path: 图像路径

        Returns:
            处理后的图像tensor
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图像文件未找到: {image_path}")

        image = Image.open(image_path).convert('RGB')
        image_tensor = self.transform(image).unsqueeze(0).to(self.device)
        return image_tensor

    def extract_watermark(self, watermarked_image):
        """
        从图像中提取水印

        Args:
            watermarked_image: 包含水印的图像tensor

        Returns:
            提取的水印消息 (decoder_C, decoder_RF)
        """
        with torch.no_grad():
            # 使用两个解码器提取水印
            decoded_message_C = self.network.encoder_decoder.module.decoder_C(watermarked_image)
            decoded_message_RF = self.network.encoder_decoder.module.decoder_RF(watermarked_image)

            return decoded_message_C, decoded_message_RF

    def calculate_error_rate(self, original_message, decoded_message):
        """
        计算解码错误率

        Args:
            original_message: 原始消息
            decoded_message: 解码消息

        Returns:
            错误率 (0-1之间，越小越好)
        """
        return self.network.decoded_message_error_rate(
            original_message.squeeze(), decoded_message.squeeze()
        )

    def compare_messages(self, original_message, decoded_message, message_type="C"):
        """
        比较原始消息和解码消息

        Args:
            original_message: 原始消息tensor
            decoded_message: 解码消息tensor
            message_type: 消息类型 ("C" 或 "RF")

        Returns:
            比较结果字典
        """
        error_rate = self.calculate_error_rate(original_message, decoded_message)

        # 转换为二进制消息进行比较
        orig_binary = (original_message.squeeze() > 0).int().cpu().numpy()
        decoded_binary = (decoded_message.squeeze() > 0).int().cpu().numpy()

        # 计算匹配的位数
        correct_bits = np.sum(orig_binary == decoded_binary)
        total_bits = len(orig_binary)
        accuracy = correct_bits / total_bits

        result = {
            'message_type': message_type,
            'error_rate': float(error_rate),  # 转换为Python float
            'accuracy': float(accuracy),  # 转换为Python float
            'correct_bits': int(correct_bits),  # 转换为Python int
            'total_bits': int(total_bits),  # 转换为Python int
            'bit_error_rate': float(1 - accuracy)  # 转换为Python float
        }

        return result

    def test_single_image(self, watermarked_image_path, original_message, output_dir=None):
        """
        测试单张图像的解码功能

        Args:
            watermarked_image_path: 水印图像路径
            original_message: 原始水印消息
            output_dir: 输出目录（可选）

        Returns:
            测试结果字典
        """
        print(f"正在测试图像: {os.path.basename(watermarked_image_path)}")

        # 加载水印图像
        watermarked_image = self.load_watermarked_image(watermarked_image_path)

        # 提取水印
        decoded_message_C, decoded_message_RF = self.extract_watermark(watermarked_image)

        # 比较结果
        result_C = self.compare_messages(original_message, decoded_message_C, "C")
        result_RF = self.compare_messages(original_message, decoded_message_RF, "RF")

        # 整合结果
        test_result = {
            'image_path': watermarked_image_path,
            'decoder_C': result_C,
            'decoder_RF': result_RF,
            'test_time': time.strftime('%Y-%m-%d %H:%M:%S')
        }

        # 打印结果
        print(f"  解码器C - 错误率: {result_C['error_rate']:.4f}, 准确率: {result_C['accuracy']:.4f}")
        print(f"  解码器RF - 错误率: {result_RF['error_rate']:.4f}, 准确率: {result_RF['accuracy']:.4f}")

        # 保存详细结果（如果指定了输出目录）
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
            image_name = os.path.splitext(os.path.basename(watermarked_image_path))[0]
            result_file = os.path.join(output_dir, f"{image_name}_decode_result.json")

            with open(result_file, 'w') as f:
                json.dump(test_result, f, indent=2)

            print(f"  详细结果保存至: {result_file}")

        return test_result

    def test_multiple_images(self, watermarked_dir, watermark_info_path, num_samples=5, output_dir=None):
        """
        测试多张图像的解码功能

        Args:
            watermarked_dir: 水印图像目录
            watermark_info_path: 水印信息文件路径
            num_samples: 测试样本数量
            output_dir: 输出目录

        Returns:
            测试结果列表
        """
        print(f"开始批量解码测试...")
        print(f"水印图像目录: {watermarked_dir}")
        print(f"测试样本数: {num_samples}")

        # 加载水印信息
        watermark_info = self.load_watermark_info(watermark_info_path)

        # 生成原始水印消息
        original_message = self.generate_original_message(13)
        print(f"原始水印消息已重新生成")

        # 获取所有水印图像文件
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
        image_files = []

        for file in os.listdir(watermarked_dir):
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(file)

        if len(image_files) == 0:
            raise ValueError(f"在 {watermarked_dir} 中未找到图像文件")

        # 随机选择测试样本
        #test_samples = random.sample(image_files, min(num_samples, len(image_files)))
        #print(f"从 {len(image_files)} 张图像中随机选择 {len(test_samples)} 张进行测试")

        # 测试每张图像
        results = []
        for i, image_file in enumerate(image_files, 1):
            print(f"\n--- 测试 {i}/{len(image_files)} ---")
            image_path = os.path.join(watermarked_dir, image_file)

            try:
                result = self.test_single_image(image_path, original_message, output_dir)
                results.append(result)
            except Exception as e:
                print(f"测试 {image_file} 时出错: {e}")
                continue

        # 生成汇总报告
        self.generate_summary_report(results, output_dir)

        return results

    def generate_summary_report(self, results, output_dir=None):
        """
        生成汇总报告

        Args:
            results: 测试结果列表
            output_dir: 输出目录
        """
        if not results:
            print("无测试结果可用于生成报告")
            return

        print(f"\n{'=' * 50}")
        print("解码测试汇总报告")
        print(f"{'=' * 50}")

        # 计算统计数据
        error_rates_C = [r['decoder_C']['error_rate'] for r in results]
        error_rates_RF = [r['decoder_RF']['error_rate'] for r in results]
        accuracies_C = [r['decoder_C']['accuracy'] for r in results]
        accuracies_RF = [r['decoder_RF']['accuracy'] for r in results]

        # 解码器C统计
        print(f"\n解码器C (通用解码器):")
        print(f"  平均错误率: {np.mean(error_rates_C):.4f}")
        print(f"  最低错误率: {np.min(error_rates_C):.4f}")
        print(f"  最高错误率: {np.max(error_rates_C):.4f}")
        print(f"  平均准确率: {np.mean(accuracies_C):.4f}")

        # 解码器RF统计
        print(f"\n解码器RF (深度伪造检测):")
        print(f"  平均错误率: {np.mean(error_rates_RF):.4f}")
        print(f"  最低错误率: {np.min(error_rates_RF):.4f}")
        print(f"  最高错误率: {np.max(error_rates_RF):.4f}")
        print(f"  平均准确率: {np.mean(accuracies_RF):.4f}")

        # 整体评估
        print(f"\n整体评估:")
        perfect_decode_C = sum(1 for rate in error_rates_C if rate < 0.01)
        perfect_decode_RF = sum(1 for rate in error_rates_RF if rate < 0.01)
        print(f"  完美解码图像数 (错误率<1%):")
        print(f"    解码器C: {perfect_decode_C}/{len(results)} ({perfect_decode_C / len(results) * 100:.1f}%)")
        print(f"    解码器RF: {perfect_decode_RF}/{len(results)} ({perfect_decode_RF / len(results) * 100:.1f}%)")

        # 保存报告
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)

            summary_report = {
                'test_summary': {
                    'total_images': len(results),
                    'test_time': time.strftime('%Y-%m-%d %H:%M:%S')
                },
                'decoder_C_stats': {
                    'avg_error_rate': float(np.mean(error_rates_C)),
                    'min_error_rate': float(np.min(error_rates_C)),
                    'max_error_rate': float(np.max(error_rates_C)),
                    'avg_accuracy': float(np.mean(accuracies_C)),
                    'perfect_decode_count': perfect_decode_C
                },
                'decoder_RF_stats': {
                    'avg_error_rate': float(np.mean(error_rates_RF)),
                    'min_error_rate': float(np.min(error_rates_RF)),
                    'max_error_rate': float(np.max(error_rates_RF)),
                    'avg_accuracy': float(np.mean(accuracies_RF)),
                    'perfect_decode_count': perfect_decode_RF
                },
                'detailed_results': results
            }

            report_file = os.path.join(output_dir, 'decode_test_summary.json')
            with open(report_file, 'w') as f:
                json.dump(summary_report, f, indent=2)

            print(f"\n详细报告已保存至: {report_file}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='SepMark水印解码测试脚本')
    parser.add_argument('--watermarked_dir', type=str, required=True,
                        help='水印图像目录路径')
    parser.add_argument('--watermark_info', type=str, required=True,
                        help='水印信息文件路径 (watermark_info.json)')
    parser.add_argument('--config', type=str, default='cfg/test_DualMark.yaml',
                        help='配置文件路径')
    parser.add_argument('--num_samples', type=int, default=100,
                        help='测试样本数量')
    parser.add_argument('--output_dir', type=str, default=None,
                        help='输出目录')
    parser.add_argument('--single_image', type=str, default=None,
                        help='测试单张图像路径')

    args = parser.parse_args()

    # 设置默认输出目录
    if args.output_dir is None:
        args.output_dir = f"decode_test_output_{int(time.time())}"

    # 创建解码器
    decoder = WatermarkDecoder(args.config)

    if args.single_image:
        # 单张图像测试
        print("单张图像解码测试模式")

        # 加载水印信息
        watermark_info = decoder.load_watermark_info(args.watermark_info)
        original_message = decoder.generate_original_message(watermark_info['watermark_seed'])

        # 测试单张图像
        result = decoder.test_single_image(args.single_image, original_message, args.output_dir)

    else:
        # 批量测试
        print("批量解码测试模式")

        # 执行批量测试
        results = decoder.test_multiple_images(
            args.watermarked_dir,
            args.watermark_info,
            args.num_samples,
            args.output_dir
        )

    print(f"\n解码测试完成！结果保存在: {args.output_dir}")


if __name__ == '__main__':
    # 如果直接运行脚本，使用固定参数
    if len(os.sys.argv) == 1:
        # 设置您的路径
        watermarked_dir = r"D:\deepfake\CelebAMask-HQ\CelebAMask-HQ\marked\watermarked"
        watermark_info_path = r"D:\deepfake\CelebAMask-HQ\CelebAMask-HQ\marked\watermark_info.json"

        print("使用默认参数进行解码测试...")
        print(f"水印图像目录: {watermarked_dir}")
        print(f"水印信息文件: {watermark_info_path}")

        # 检查路径是否存在
        if not os.path.exists(watermarked_dir):
            print(f"错误: 水印图像目录不存在: {watermarked_dir}")
            exit(1)

        if not os.path.exists(watermark_info_path):
            print(f"错误: 水印信息文件不存在: {watermark_info_path}")
            exit(1)

        # 创建解码器并执行测试
        decoder = WatermarkDecoder()
        results = decoder.test_multiple_images(
            watermarked_dir,
            watermark_info_path,
            num_samples=3,  # 测试3张图片
            output_dir=f"decode_test_output_{int(time.time())}"
        )
    else:
        main()
